#!/usr/bin/python
# -*-coding:utf-8-*-
import os
import pandas as pd
import numpy as np

### 底层读取数据的依赖（不提供）
from zbc_factor_lib.base.factors_library_base import NewRQFactorLib as DataReader

from dir_info import ylabel_data_dir

data_reader = DataReader()
# data_reader.show_basic_library_db()

def get_n_day_return_label_main(n_days=5):
    # TODO - 读取数据
    processed_daily_stock_trade_data = data_reader.read_basic_data_table('processed_daily_stock_trade_data')

    processed_daily_stock_trade_data = processed_daily_stock_trade_data.rename(columns={'code': 'stock_code'})

    processed_daily_stock_trade_data['bclose'] = \
        processed_daily_stock_trade_data['close'] * processed_daily_stock_trade_data['back_adjfactor']

    processed_daily_stock_trade_data = processed_daily_stock_trade_data.sort_values(['stock_code', 'date'], ascending=True)

    processed_daily_stock_trade_data = processed_daily_stock_trade_data.set_index(['date', 'stock_code'])

    y_label = processed_daily_stock_trade_data.groupby('stock_code')[['bclose']].pct_change(-n_days)

    y_label = y_label.rename(columns={'bclose': 'y'})

    # print(y_label.head())
    #
    # # TODO - 映射
    # y_label = y_label.reset_index()
    #
    # y_label = trade_date_map(y_label, key='date', reverse=False)
    # y_label = stock_code_map(y_label, key='stock_code', reverse=False)
    #
    # y_label = y_label.set_index(['date', 'stock_code'])

    y_label['y'] = y_label['y'].astype(np.float32)

    # print(y_label.head())

    # TODO - save
    save_filename = 'y_%dday_return_label' % (n_days)

    y_label.to_hdf(os.path.join(ylabel_data_dir, save_filename+'.h5'), key=save_filename)

    print(save_filename, 'done!\n')


if __name__ == '__main__':
    n_days_list = [3, 5, 10, 20]

    for n_days in n_days_list:
        get_n_day_return_label_main(n_days=n_days)

